#! /usr/bin/python
# -*- encoding: utf-8 -*-

import torch
import torch.nn as nn
import torch.nn.functional as F
from .utils import accuracy


class LossFunction(nn.Module):
	def __init__(self, nOut, nClasses, **kwargs):
		super(LossFunction, self).__init__()
		self.fc = nn.Linear(nOut, nClasses)
		self.criertion = nn.CrossEntropyLoss()

	def forward(self, x, label=None):
		x = self.fc(x)
		loss = self.criertion(x, label)
		prec1 = accuracy(x.detach(), label.detach(), topk=(1,))[0]
		return loss, prec1

